package com.enhanced.mqtt.client;

import org.eclipse.paho.client.mqttv3.*;
import org.eclipse.paho.client.mqttv3.persist.MemoryPersistence;

import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManagerFactory;
import java.io.FileInputStream;
import java.io.InputStream;
import java.security.KeyStore;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.Scanner;

/**
 * SSL/TLS MQTT客户端示例
 * 演示如何使用SSL/TLS加密连接到EMQX
 */
public class SSLMqttClient {
    
    private static final String SSL_BROKER_URL = "ssl://192.168.3.234:8883";
    private static final String TOPIC = "test/ssl";
    private static final int QOS = 1;
    
    public static void main(String[] args) {
        try {
            runDemo(new Scanner(System.in));
        } catch (Exception e) {
            System.out.println("MQTT错误: " + e.getMessage());
            e.printStackTrace();
        }
    }
    
    /**
     * 运行演示，使用外部传入的Scanner
     * @param scanner 已经打开的Scanner
     */
    public static void runDemo(Scanner scanner) throws Exception {
        // 创建SSL连接示例
        System.out.println("=== SSL/TLS加密连接示例 ===");
        System.out.println("本演示展示如何使用SSL/TLS加密连接到EMQX服务器");
        
        // 提示用户输入服务器信息和证书路径
        System.out.print("是否使用CA证书验证？(y/n): ");
        String useCA = scanner.nextLine().trim().toLowerCase();
        
        MqttConnectOptions options = new MqttConnectOptions();
        options.setCleanSession(true);
        
        if ("y".equals(useCA)) {
            System.out.print("请输入CA证书路径: ");
            String caFile = scanner.nextLine().trim();
            
            // 配置SSL上下文
            SSLSocketFactory socketFactory = getSocketFactory(caFile);
            options.setSocketFactory(socketFactory);
            System.out.println("已配置CA证书验证");
        } else {
            // 不验证证书
            options.setSocketFactory(getSocketFactoryWithoutVerification());
            System.out.println("已配置不验证证书的SSL连接");
        }
        
        // 创建客户端
        String clientId = "SSLClient-" + System.currentTimeMillis();
        MqttAsyncClient client = new MqttAsyncClient(SSL_BROKER_URL, clientId, new MemoryPersistence());
        
        // 设置回调
        client.setCallback(new MqttCallback() {
            @Override
            public void connectionLost(Throwable cause) {
                System.out.println("连接丢失: " + cause.getMessage());
            }

            @Override
            public void messageArrived(String topic, MqttMessage message) {
                String payload = new String(message.getPayload());
                System.out.println("\n收到消息:");
                System.out.println("  主题: " + topic);
                System.out.println("  内容: " + payload);
                System.out.println("  QoS: " + message.getQos());
            }

            @Override
            public void deliveryComplete(IMqttDeliveryToken token) {
                System.out.println("消息发送完成");
            }
        });
        
        // 连接服务器
        System.out.println("正在连接到SSL MQTT服务器: " + SSL_BROKER_URL);
        client.connect(options).waitForCompletion(5000);
        System.out.println("已连接到SSL MQTT服务器，客户端ID: " + clientId);
        
        // 订阅主题
        client.subscribe(TOPIC, QOS).waitForCompletion();
        System.out.println("已订阅主题: " + TOPIC);
        
        // 发送消息
        String message = "这是一条通过SSL/TLS加密的消息";
        MqttMessage mqttMessage = new MqttMessage(message.getBytes());
        mqttMessage.setQos(QOS);
        client.publish(TOPIC, mqttMessage).waitForCompletion();
        System.out.println("已发送加密消息");
        
        // 等待接收消息
        System.out.println("等待5秒接收消息...");
        Thread.sleep(5000);
        
        // 断开连接
        client.disconnect().waitForCompletion();
        System.out.println("已断开SSL连接");
        
        // 不关闭scanner，由调用者负责关闭
    }
    
    /**
     * 加载CA证书并创建SSLSocketFactory
     */
    private static SSLSocketFactory getSocketFactory(String caCertFile) throws Exception {
        // 加载CA证书
        CertificateFactory cf = CertificateFactory.getInstance("X.509");
        X509Certificate caCert;
        
        try (InputStream caStream = new FileInputStream(caCertFile)) {
            caCert = (X509Certificate) cf.generateCertificate(caStream);
        }
        
        // 创建信任库并加载CA证书
        KeyStore caKs = KeyStore.getInstance(KeyStore.getDefaultType());
        caKs.load(null, null);
        caKs.setCertificateEntry("ca-certificate", caCert);
        
        // 创建信任管理器
        TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
        tmf.init(caKs);
        
        // 创建SSL上下文
        SSLContext sslContext = SSLContext.getInstance("TLSv1.2");
        sslContext.init(null, tmf.getTrustManagers(), null);
        
        return sslContext.getSocketFactory();
    }
    
    /**
     * 创建不验证证书的SSLSocketFactory
     */
    public static SSLSocketFactory getSocketFactoryWithoutVerification() throws Exception {
        SSLContext sslContext = SSLContext.getInstance("TLSv1.2");
        sslContext.init(null, null, null);
        return sslContext.getSocketFactory();
    }
} 